use alloc::alloc::Layout;
use alloc::alloc::{alloc as std_alloc, dealloc as std_dealloc};
use alloc::vec::Vec;
use core::mem::align_of;
use core::ops::{Deref, DerefMut, Index, IndexMut};

#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Hash)]
pub struct Page {
    ptr: *mut u8,
    size: usize,
}

impl Page {
    pub fn new(initial_size: usize) -> Self {
        let layout = Layout::from_size_align(initial_size, align_of::<u8>()).unwrap();
        unsafe {
            let ptr = std_alloc(layout);
            if ptr.is_null() {
                panic!("Memory allocation failed");
            }
            Page {
                ptr,
                size: initial_size,
            }
        }
    }

    pub fn set_size(&mut self, new_size: usize) -> Result<(), ()> {
        if new_size == self.size {
            return Ok(());
        }

        let layout = Layout::from_size_align(self.size, align_of::<u8>()).unwrap();
        unsafe {
            let new_ptr = std_alloc(Layout::from_size_align(new_size, align_of::<u8>()).unwrap());
            if new_ptr.is_null() {
                return Err(());
            }

            core::ptr::copy_nonoverlapping(self.ptr, new_ptr, self.size.min(new_size));

            std_dealloc(self.ptr, layout);

            self.ptr = new_ptr;
            self.size = new_size;
            Ok(())
        }
    }

    pub fn grow(&mut self, new_size: usize) -> Result<(), ()> {
        let old_size = self.size();
        if let Some(size) =
            old_size.checked_mul(usize::max((new_size / old_size + 1).next_power_of_two(), 1))
        {
            self.set_size(size)
        } else {
            Err(())
        }
    }

    pub fn as_mut_ptr(&mut self) -> *mut u8 {
        self.ptr
    }

    pub fn size(&self) -> usize {
        self.size
    }
}

impl Deref for Page {
    type Target = [u8];

    fn deref(&self) -> &Self::Target {
        unsafe { core::slice::from_raw_parts(self.ptr, self.size) }
    }
}

impl DerefMut for Page {
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe { core::slice::from_raw_parts_mut(self.ptr, self.size) }
    }
}

pub struct Pages(pub Vec<Page>);

impl Pages {
    pub fn new(page_num: usize, initial_size: usize) -> Self {
        let mut pages = Vec::with_capacity(page_num);
        for _ in 0..page_num {
            pages.push(Page::new(initial_size));
        }
        Pages(pages)
    }

    pub fn grow(&mut self, new_size: usize) -> Result<(), ()> {
        for page in self.0.iter_mut() {
            page.grow(new_size)?;
        }
        Ok(())
    }

    pub fn size(&self) -> usize {
        if self.0.is_empty() {
            0
        } else {
            self.0[0].size()
        }
    }
}

impl Index<usize> for Pages {
    type Output = Page;

    fn index(&self, index: usize) -> &Self::Output {
        &self.0[index]
    }
}

impl IndexMut<usize> for Pages {
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        &mut self.0[index]
    }
}
